use DBI;

package ExploitFarm_SQL;
use sql_query;
use Util::Logger;

###sql.pl -- our own API to work with PostgreSQL.

sub new
{
    my ($self,$db_host,$db_port,$db_name,$db_login,$db_pass)=@_;
    @_=grep (/^Pg$/,DBI->available_drivers);
    die ("FATAL: No DBI driver for PostgreSQL (dbi:pg)!") if ($#_<0);
    my %ret=(
	    db_host	=> $db_host,
	    db_port	=> $db_port,
	    db_name	=> $db_name,
	    db_login	=> $db_login,
	    db_pass	=> $db_pass,
	);
    my $l=new Util::Logger('getters/sql_log.txt') or die "SQL: Can't open log file ! : $!";
    $ret{'log'}=$l;
    bless \%ret;
}


sub connect
{
    my $self=shift;
    my $ret=0;
    my $db_name=$self->{db_name};
    my $db_host=$self->{db_host};
    my $db_port=$self->{db_port};
    my $db_login=$self->{db_login};
    my $db_pass=$self->{db_pass};

    my $dbh = DBI->connect_cached("dbi:Pg:dbname=$db_name;host=$db_host;port=$db_port;" .
                          "", "$db_login", "$db_pass") or $ret="$DBI::errstr";
    $dbh->{AutoCommit} = 0;
    $dbh->{RaiseError} = 1;
    $self->{dbh}=$dbh;
    return $ret if ($ret);
    return undef;
}

sub disconnect
{
    my $self=shift;
    my $dbh=$self->{dbh};
    $dbh->disconnect();
}

sub prepare
{
    my ($self,$str,$sql_params)=@_;
    my $dbh=$self->{'dbh'};
    $self->{'log'}->debug("PREPARing0: $str");
    if (!exists $sqlq{$str})
    {
	die "'$str':No such sql id!";
    }
    $str = $sqlq{$str};
    $str =~ s/\!\!(\w+?)\!\!/$$sql_params{$1}/g;
    $self->{'log'}->debug("PREPARing1: $str");
    $sth = $dbh->prepare($str);
    $self->{'log'}->info("PREPARE: $str");
    $self->{dbh}=$dbh;
    $self->{sth}=$sth;
    return undef;
}

sub exec
{
    my ($self,$sql_str,$sql_params)=@_;
    my $dbh=$self->{dbh};
    my $sth;
    my $ret=undef;
    $self->{'log'}->debug("EXEC0: $sql_str");
    for (1..3)
    {
	eval #try ... catch
	{
	    $self->prepare($sql_str,$sql_params);
	    $sth=$self->{sth};
	    $sth->execute();
	    last;
	};
	if ($@ =~ /duplicate key value violates unique constraint/c)
	{
	    last;
	    eval
	    {
		$sth->rollback() if (defined $sth);
	    };
	}
	elsif($@)
	{
	    $self->{'log'}->debug("EXEC_EXEPTIONS: $@");
	    $sth=$self->{sth};
	    eval
	    {
		$sth->rollback() if (defined $sth);
	    };
	    $ret=$@;
	    $self->connect();
	}
    }
    $self->{dbh}=$dbh;
    $self->{sth}=$sth;
    $self->{'log'}->info("EXEC: $sql_str");
    return $ret;
}

sub exec_many
{
    my ($self,$sql_str,$sql_params)=@_;
    my $dbh=$self->{dbh};
    my $sth;
    my $ret=undef;
    $self->{'log'}->debug("EXEC0: $sql_str");
    for (1..3)
    {
	eval #try ... catch
	{
	    for (@$sql_str)
	    {
		$self->prepare($_,$sql_params);
		$sth=$self->{sth};
		$sth->execute();
		last;
	    }
	};
	if ($@ =~ /duplicate key value violates unique constraint/c)
	{
	    last;
	    eval
	    {
		$sth->rollback() if (defined $sth);
	    };
	}
	elsif($@)
	{
	    $self->{'log'}->debug("EXEC_EXEPTIONS: $@");
	    $sth=$self->{sth};
	    eval
	    {
		$sth->rollback() if (defined $sth);
	    };
	    $ret=$@;
	    $self->connect();
	}
    }
    $self->{dbh}=$dbh;
    $self->{sth}=$sth;
    $self->{'log'}->info("EXEC: $sql_str");
    return $ret;
}

sub fetchrow_array
{
    my $self=shift;
    my $sth=$self->{sth};
    return $sth->fetchrow_array;
}

sub commit
{
    my $self=shift;
    my $dbh=$self->{dbh};
    $self->{'log'}->info("COMMIT");
    return $dbh->commit;
}

sub rollback
{
    my $self=shift;
    my $dbh=$self->{dbh};
    $self->{'log'}->info("ROLLBACK");
    return $dbh->rollback;
}

sub get_teams
{
    my $self=shift;
    my $dbh=$self->{dbh};
    my $sql_exec;
    my %conf;
    my @ret;
    $sql_exec=$self->exec('get_teams',\%conf);
    return [undef,$sql_exec] if (defined $sql_exec);
    while (@row=$self->fetchrow_array)
    {
	push(@ret,$row[0]);
    }
    return \@ret;
}

sub set_teams
{
    my ($self,$teams)=@_;
    my $sql_exec;
    my %conf;
    my $teams_values='';
    $sql_exec=$self->exec('trunk_teams',\%conf);
    if (defined $sql_exec)
    {
	return $sql_exec;
    }
    for (keys %$teams)
    {
	$teams_values.="('$_','$$teams{$_}',true),";
    }
    substr($teams_values,length($teams_values)-1)='';
    $conf{'teams_values'}=$teams_values;
    $sql_exec=$self->exec('set_teams',\%conf);
    if (defined $sql_exec)
    {
	$self->rollback();
	return $sql_exec;
    }
    $self->commit();
    return undef;
}

sub put_flags
{
    my ($self,$exploit,$ip,$flags)=@_;
    my %conf;
    my $sql_exec;

    $conf{'team'}=$ip;
    $conf{'exploit'}=$exploit;
    for (@$flags)
    {
	$conf{'flag'}=$_;
	$sql_exec=$self->exec('put_flag',\%conf);
    }
    $self->commit();
}

sub get_flags
{
    my $self=shift;

    my %conf;
    my $sql_exec;
    my $ret={};

    $sql_exec=$self->exec('get_flags',\%conf);
    if (defined $sql_exec)
    {
	$self->rollback();
	return {error => $sql_exec};
    }

    while (@row=$self->fetchrow_array)
    {
	$$ret{$row[1]}=$row[0];
    }

    my @flags=values %$ret;
    $self->update_fstatus(1,\@flags);

    return $ret;
}

sub update_fstatus
{
    my ($self,$status,$flags)=@_;
    my %conf;
    my $sql_exec;

    my $str='';
    for (@$flags)
    {
	$str.=" fid='$_' OR";
    }
    substr($str,length($str)-2)='  ';

    $conf{fids}=$str;
    $conf{status}=$status;
    $sql_exec=$self->exec('update_flags_fstatus',\%conf);
    return $sql_exec if (defined $sql_exec);

    $self->commit();

    return $sql_exec;
}

sub clean_fstatus_1
{
    my $self=shift;
    my %conf;
    my $sql_exec;

    $sql_exec=$self->exec('clean_fstatus_1',\%conf);
    return $sql_exec if (defined $sql_exec);

    $self->commit();
}

1;
